import os
import shutil

if os.path.isdir('mydata'):
    shutil.rmtree('mydata')
os.makedirs('mydata/train/img')
os.makedirs('mydata/train/xml')
os.makedirs('mydata/eval/img')
os.makedirs('mydata/eval/xml')

label_list=open('mydata/label_list.txt','w')
train_list=open('mydata/train_list.txt','w')
eval_list=open('mydata/eval_list.txt','w')

for root,dirs,files in os.walk('dataset_origin/'):
    for dir in dirs:
        label=dir
        label_list.write(label + '\n')
        print('Find class \"' + label + '\"')
        train_path='dataset_origin/' + dir + '/train/'
        eval_path='dataset_origin/' + dir + '/eval/'
        for root2,dirs2,files2 in os.walk(train_path + 'img'):
            print('train of \"' + label + '\" found ' + str(len(files2)))
            for img_name in files2:
                xml_name=img_name.replace('.jpg','.xml')
                shutil.copy(train_path + 'img/' + img_name,'mydata/train/img/' + img_name.replace(' ','_'))
                shutil.copy(train_path + 'xml/' + xml_name,'mydata/train/xml/' + xml_name.replace(' ','_'))
                train_list.write('train/img/' + img_name.replace(' ','_') + ' train/xml/' + xml_name.replace(' ','_') + '\n')
        for root2,dirs2,files2 in os.walk(eval_path + 'img'):
            print('eval of \"' + label + '\" found ' + str(len(files2)))
            for img_name in files2:
                xml_name=img_name.replace('.jpg','.xml')
                shutil.copy(eval_path + 'img/' + img_name,'mydata/eval/img/' + img_name.replace(' ','_'))
                shutil.copy(eval_path + 'xml/' + xml_name,'mydata/eval/xml/' + xml_name.replace(' ','_'))
                eval_list.write('eval/img/' + img_name.replace(' ','_') + ' eval/xml/' + xml_name.replace(' ','_') + '\n')
        print()
    break
print('Compelete!')
